# core
import draw

from importlib.metadata import entry_points
from future.utils import viewitems

from miasm.arch.x86.regs import *
from miasm.expression.expression import  ExprInt, ExprOp
from miasm.analysis.simplifier import IRCFGSimplifierCommon, IRCFGSimplifierSSA
from miasm.core.asmblock import AsmConstraint,AsmCFG
from miasm.expression.simplifications import expr_simp
from mod_utils import *
from miasm.analysis.cst_propag import propagate_cst_expr
from miasm.analysis.data_flow import DeadRemoval,remove_empty_assignblks

from z3 import *
from mylib import stop_on_jmp,ActionList
from BaseSolver import BaseSolver

def get_backbone_blks(asmcfg:AsmCFG):
    blks_list=[]
    for blk in asmcfg.blocks:
        if blk.lines[0].name=='CMP' or blk.lines[0].name=='JZ':
            offset=asmcfg.loc_db.get_location_offset(blk.loc_key)
            blks_list.append(offset)
    return blks_list

def fla_ollvm(asmcfg,head,ircfg,lifter,mdis,loc_db):
    action=ActionList()
    mdis.dis_block_callback=stop_on_jmp
    head_blk=asmcfg.loc_key_to_block(head)
    head_off=head_blk.get_offsets()[0]
    main_asmcfg=asmcfg
    main_ircfg=ircfg
    
    relevant_blocks, dispatcher, pre_dispatcher = get_cff_info(main_asmcfg)
    dispatcher_blk = main_asmcfg.getby_offset(dispatcher)
    # 分发器的第一条命令
    dispatcher_first_instr = dispatcher_blk.lines[0]
    # 状态变量:
    state_var = dispatcher_first_instr.get_args_expr()[1]
    print(state_var)
    backbone = {}
    fixed_cfg = {}
    val_list = []
    rel_blk_info = {}
    # mdis.dis_block_callback = stop_on_jmp
    for addr in relevant_blocks:
        print("Getting info for relevant block @ %#x"%addr)
        asmcfg = mdis.dis_multiblock(addr)
        ircfg = lifter.new_ircfg_from_asmcfg(asmcfg)
        ircfg_simplifier = IRCFGSimplifierCommon(lifter)
        ircfg_simplifier.simplify(ircfg, addr)

        # 标记影响状态变量的指令为nop_addrs
        # marking the instructions affecting the state variable as nop_addrs
        nop_addrs = find_state_var_usedefs(ircfg, state_var)
        rel_blk_info[addr] = (asmcfg, nop_addrs)
        # print(hex(addr),nop_addrs)

        # 返回offset对应的LocKey
        head = loc_db.get_offset_location(addr)
        # 将原来的IR转化为SSA
        ssa_simplifier = IRCFGSimplifierSSA(lifter)
        ssa = ssa_simplifier.ircfg_to_ssa(ircfg, head)
        #we only use do_propagate_expressions ssa simp pass
        ssa_simplifier.do_propagate_expressions(ssa, head)

        # 得到可能的状态值
        var_asg, tmpval_list = find_var_asg(ircfg, {state_var})
        # print(var_asg,tmpval_list)
        print('%#x %s' % (addr, var_asg))

        # adding all the possible values to a global list
        val_list += tmpval_list

        # 最后一个基本块
        last_blk = list(asmcfg.blocks)[-1]
        # checking the type of relevant blocks on the basis of no. of possible values
        if len(var_asg) == 1:
            # var_asg['next'] = hex(var_asg['next'])
            #map value of state variable in rel block
            fixed_cfg[addr] = var_asg
        # 存在条件跳转
        elif len(var_asg) > 1:
            #extracting the condition from the last 3rd line
            # print(last_blk.lines[-3])
            cond_mnem = last_blk.lines[-3].name
            print('cond used: %s' % cond_mnem)
            var_asg['cond'] = cond_mnem
            # var_asg['true_next'] = hex(var_asg['true_next'])
            # var_asg['false_next'] = hex(var_asg['false_next'])
            # map the conditions and possible values dictionary to the cfg info
            fixed_cfg[addr] = var_asg
        elif len(last_blk.lines)==1 and len(var_asg)==0:
                #tail has a single instruction ie. jmp and no assignments
                tail = addr
                print("found backbone tail @ %#x" % addr)
        else:
            print("no state variable assignments found for relevant block @ %#x" % addr)
            # return empty patches as deobfuscation failed!!
            #return {}
    print('val_list: ' + ', '.join([hex(val) for val in val_list]))
    
    for lbl, irblock in viewitems(main_ircfg.blocks):
        # lbl:loc_key,irblock:{assignblks}
        for assignblk in irblock:
            # assignblk:一行语句
            asg_items = assignblk.items()# 转换为IR表示
            if asg_items:    # do not enter if nop
                dst, src = asg_items[0] #IR表示的一条语句
                if isinstance(src, ExprOp):
                    if src.op == 'FLAG_EQ_CMP': # 捕获比较操作
                        arg = src.args[1]   #src:{op,args}
                        if isinstance(arg, ExprInt):
                            if int(arg) in val_list:    # 判断这个比较是否为对于状态值的比较，我推测寻找backbone
                                cmp_val = int(arg)
                                # 猜测当前判断这个块是在进行状态值为cmp_val时的流程分发，最后一条指令应该为jcc
                                # 跳转的IR有两条式子，对EIP和IRDst的操作
                                # var应该为EIP，而locs为判断条件，第一个跳转方向为真实的方向
                                var, locs = irblock[-1].items()[0]
                                # 真正的跳转方向
                                true_dst = main_ircfg.loc_db.get_location_offset(locs.src1.loc_key)
                                # 记录map，在比较值为cmp_val时，分发的块地址为true_dst
                                backbone[cmp_val] = true_dst
    fixed_cfg[0x80484C0]={'next':0x885336A4}
    for offset, link in fixed_cfg.items():
        if 'cond' in link:  # 如果存在跳转
            tval = fixed_cfg[offset]['true_next']
            fval = fixed_cfg[offset]['false_next']
            fixed_cfg[offset]['true_next'] = backbone[tval]
            fixed_cfg[offset]['false_next'] = backbone[fval]
        elif 'next' in link:
            fixed_cfg[offset]['next'] = backbone[link['next']]
        else:
            # the tail doesn't has any condition
            tail = int(offset, 16)
    rel_blk_info.pop(tail)  
    print('fixed_cfg',fixed_cfg)
    
    
    
    tail = main_asmcfg.getby_offset(tail).lines[-1]
    
    
    
    backbone_start, backbone_end = dispatcher, tail.offset + tail.l
    print('backbone_start',hex(backbone_start))
    print('backbone_end',hex(backbone_end))
    
    # 修改asmcfg
    asmcfg_cut=main_asmcfg.copy()
    # 删除所有的backbone
    backbone_list=get_backbone_blks(main_asmcfg)
    for blk in main_asmcfg.blocks:
        blk_off=blk.get_offsets()[0]
        if blk_off==pre_dispatcher:
            action.graph_del_blk(blk_off)
            asmcfg_cut.del_block(asmcfg_cut.getby_offset(blk_off))
        if blk_off>=backbone_start and blk_off<backbone_end:
            action.graph_del_blk(blk_off)
            asmcfg_cut.del_block(asmcfg_cut.getby_offset(blk_off))
    print('pre_dispatcher',hex(pre_dispatcher))
    print(backbone)
    for fixed_itme in fixed_cfg:
        src_off=fixed_itme
        src_loc=asmcfg_cut.getby_offset(src_off).loc_key
        dst_item=fixed_cfg[src_off]
        
        if 'cond' in dst_item:
            dst_true=dst_item['true_next']
            dst_true_loc=asmcfg_cut.getby_offset(dst_true).loc_key
            asmcfg_cut.add_edge(src_loc,dst_true_loc,AsmConstraint(dst_true_loc,'c_to'))
            action.graph_add_edge(src_off,dst_true)
            dst_false=dst_item['false_next']
            dst_false_loc=asmcfg_cut.getby_offset(dst_false).loc_key
            asmcfg_cut.add_edge(src_loc,dst_false_loc,AsmConstraint(dst_false_loc,'c_to'))
            action.graph_add_edge(src_off,dst_false)
        if 'next' in dst_item:
            dst_next=dst_item['next']
            dst_next_loc=asmcfg_cut.getby_offset(dst_next).loc_key
            asmcfg_cut.add_edge(src_loc,dst_next_loc,AsmConstraint(dst_next_loc,'c_to'))
            action.graph_add_edge(src_off,dst_next)
    draw.draw_graph(asmcfg_cut,"asmcfg")
    return action

def check_path_feasibility(conds):
    solver = z3.Solver()
    # 遍历每个cond，lval为cond_expr，右侧为cond_expr值
    for lval, rval in conds:
        z3_cond = Translator.to_language("z3").from_expr(lval)
        solver.add(z3_cond == int(rval.arg))

    rslt = solver.check()

    if rslt == z3.sat:
        return True
    else:
        return False
class FinalState:
    def __init__(self, result, sym, path_conds, path_history):
        self.result = result
        self.sb = sym
        self.path_conds = path_conds
        self.path_history = path_history

def explore(ir, start_addr, start_symbols, 
        ircfg, cond_limit=30, uncond_limit=100, 
        lbl_stop=None, final_states=[]):

    def codepath_walk(addr, symbols, conds, depth, final_states, path):

        if depth >= cond_limit:
            warnings.warn("'depth' is over the cond_limit :%d"%(depth))
            return 

        sb = SymbolicExecutionEngine(ir, symbols)

        for _ in range(uncond_limit):
            if isinstance(addr, ExprInt): 
                if int(addr) == lbl_stop:
                    final_states.append(FinalState(True, sb, conds, path))
                    return

            path.append(addr)

            pc = sb.run_block_at(ircfg, addr)

            if isinstance(pc, ExprCond): 
    
                # Calc the condition to take true or false paths
                cond_true  = {pc.cond: ExprInt(1, 32)}
                cond_false = {pc.cond: ExprInt(0, 32)}

                # The destination addr of the true or false paths
                addr_true  = expr_simp(
                    sb.eval_expr(pc.replace_expr(cond_true), {}))

                addr_false = expr_simp(
                    sb.eval_expr(pc.replace_expr(cond_false), {}))

                # Need to add the path conditions to reach this point
                conds_true = list(conds) + list(cond_true.items())
                conds_false = list(conds) + list(cond_false.items())

                # TODO:
                # Please add some code to complete the SMT-based path explore
                # Hint1: use check_path_feasibility()
                # Hint2: Do not forget to add the current state to the final_states 
                # when the path condition is infeasible. 
                # e.g., final_states.append(FinalState(False, sb, conds_true/conds_false, path))
                #
                # From here --------------------------------
                sat_true=check_path_feasibility(conds_true)
                sat_false=check_path_feasibility(conds_false)
                if sat_true:
                    codepath_walk(
                            addr_true, sb.symbols.copy(), 
                            conds_true, depth + 1, final_states, list(path))
                else:
                    final_states.append(FinalState(False, sb, conds_true, path))
                if sat_false:
                    codepath_walk(
                            addr_false, sb.symbols.copy(), 
                            conds_false, depth + 1, final_states, list(path))
                else:
                    final_states.append(FinalState(False, sb, conds_false, path))

                # To here --------------------------------

                return
            else:
                addr = expr_simp(sb.eval_expr(pc))

        final_states.append(FinalState(True, sb, conds, path))
        return 

    return codepath_walk(start_addr, start_symbols, [], 0, final_states, [])

def bcf_ollvm(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db):
    action=ActionList()
    mdis.dis_block_callback=stop_on_jmp
    func_addr=func_now.start_ea
    ir_arch=machine.ira(loc_db)
    symbols_init =  {
        # ExprMem(ExprId('ESP_init', 32), 32) : ExprInt(0xdeadbeef, 32)
    }

    for i, r in enumerate(all_regs_ids):
        symbols_init[r] = all_regs_ids_init[i]
    symbols_init[ExprMem(ExprId('ESP_init', 32), 32)] = ExprInt(0xdeadbeef, 32)
    final_states = []
    explore(ir_arch, 
        func_addr, 
        symbols_init, 
        ircfg, 
        # lbl_stop=0x8048D51,
        lbl_stop=0xdeadbeef, 
        final_states=final_states)
    print('final states:', len(final_states))
    counter=0
    for final_state in final_states:
        if final_state.result:
            counter+=1
    print(counter)
    executed_lockey   = []
    unexecuted_lockey = []
    unexecuted_offset=[]
    # The IR nodes which are included in one of paths were executed.
    for final_state in final_states:
        if final_state.result:
            for node in final_state.path_history:
                if isinstance(node, int):
                    lbl = ircfg.get_loc_key(node)
                elif isinstance(node, ExprInt):
                    lbl = ircfg.get_loc_key(node)
                elif isinstance(node, LocKey):
                    lbl = node.loc_key

                if lbl not in executed_lockey:
                    executed_lockey.append(lbl)

    # Otherwise, the IR nodes which are not included in any path were not executed.
    for lbl, irblock in viewitems(ircfg.blocks):
        if lbl not in executed_lockey:
            unexecuted_lockey.append(lbl)
            unexecuted_off=asmcfg.loc_key_to_block(lbl).get_offsets()[0]
            unexecuted_offset.append(hex(unexecuted_off))
    print(unexecuted_lockey)
    print(unexecuted_offset)
    asmcfg_copy=asmcfg.copy()
    for del_loc in unexecuted_lockey:
        del_blk=asmcfg_copy.loc_key_to_block(del_loc)
        action.graph_del_blk(del_blk.get_offsets()[0])
        asmcfg_copy.del_block(del_blk)
    
    draw.draw_graph(asmcfg_copy,'asmcfg')
    return action

def sub_ollvm(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db):
    mdis.dis_block_callback=stop_on_jmp
    func_addr=func_now.start_ea
    ir_arch=machine.ira(loc_db)
    init_infos = ir_arch.arch.regs.regs_init
    cst_propag_link = propagate_cst_expr(ir_arch, ircfg, func_addr, init_infos)
    deadrm = DeadRemoval(ir_arch)
    
    modified = True
    while modified:
        modified = False
        modified |= deadrm(ircfg)
        modified |= remove_empty_assignblks(ircfg)
    draw.draw_graph(ircfg,'ircfg')
    return
class OllvmFlaSolver(BaseSolver):
    def __init__(self) -> None:
        super().__init__('Ollvm-Fla','fla')
        pass
    
    def Solve(self,obj):
        super().Solve(obj)
        self.lift()
        self.action=fla_ollvm(self.asmcfg,self.head,self.ircfg,self.lifter,self.mdis,self.loc_db)
        return
class OllvmBcfSolver(BaseSolver):
    def __init__(self) -> None:
        super().__init__('Ollvm-Bcf','bcf')
        pass
    
    def Solve(self,obj):
        super().Solve(obj)
        self.lift()
        self.action=bcf_ollvm(self.func_now,self.machine,self.asmcfg,self.head,self.ircfg,self.lifter,self.mdis,self.loc_db)
        return

if __name__=='__main__':
    print('import test')
    
def test_func(index):
    print('let me see see!')

